from typing import Any
from bson import ObjectId
from pydantic_core import core_schema
from pydantic import GetCoreSchemaHandler


class PyObjectId(ObjectId):
    @classmethod
    def __get_pydantic_core_schema__(
            cls,
            source_type: Any,
            handler: GetCoreSchemaHandler,
    ) -> core_schema.CoreSchema:
        # 验证逻辑：确保输入值是有效的 ObjectId
        def validate_from_str(v: str) -> ObjectId:
            if not ObjectId.is_valid(v):
                raise ValueError("Invalid ObjectId")
            return ObjectId(v)

        # 定义一个核心模式，它首先尝试将输入作为 ObjectId 实例处理，
        # 如果失败，则尝试将其作为字符串并通过我们的验证函数处理。
        # 这是处理来自 Python 代码的数据（可以是 str 或 ObjectId）
        python_schema = core_schema.union_schema(
            [
                core_schema.is_instance_schema(ObjectId),
                core_schema.no_info_plain_validator_function(validate_from_str),
            ],
            serialization=core_schema.plain_serializer_function_ser_schema(lambda x: str(x)),
        )

        # 定义当数据来自 JSON 时的处理方式。
        # JSON 中的 ObjectId 总是字符串。
        # 我们告诉 Pydantic，JSON Schema 中这个字段的类型是 'string'。
        # 验证逻辑仍然使用我们的自定义函数。
        # 序列化逻辑将 ObjectId 转换为字符串。
        return core_schema.json_or_python_schema(
            python_schema=python_schema,
            json_schema=core_schema.str_schema(),
        )